#pragma once
#include <iostream>
using namespace std;
#include <set>
#include <map>
#include <utility>
#include <assert.h>

template <class K, class V>
struct AVLTreeNode
{
    typedef AVLTreeNode<K,V> Self;

    Self* _left;
    Self* _right;
    Self* _parent;
    int _bf;//
    pair<K,V> _kv;

    AVLTreeNode(const pair<K,V>& kv)
    :_kv(kv),_left(nullptr),_right(nullptr),_parent(nullptr),_bf(0)
    {}
};

template <class K, class V>
class AVLTree
{
    typedef AVLTreeNode<K,V> Node;
public:
    AVLTree() : _root(nullptr) {}
    bool insert(const pair<K,V>& kv);

    void inorder()
    {
        _inorder(_root);
    }

    int size()
    {
        return _size(_root);
    }

    int height()
    {
        return _height(_root);
    }

    bool isbalance()
    {
        int height = 0;
        return _isbalance(_root,height);
    }

private:
    void RotateL(Node* parent);
    void RotateR(Node* parent);
    void RotateLR(Node* parent);
    void RotateRL(Node* parent);
    void _inorder(Node* root);
    int _height(Node* root);
    bool _isbalance(Node* root,int& height);
    int _size(Node* root);
private:
    Node* _root;
};

template <class K, class V>
int AVLTree<K,V>::_size(Node* root)
{
    if(root == nullptr)
    {
        return 0;
    }
    
    return _size(root->_left) + _size(root->_right) + 1;
}

template <class K, class V>
bool AVLTree<K,V>::_isbalance(Node* root,int& height)
{
    if(root == nullptr)
    {
        height = 0;
        return true;
    }

    int leftheight = 0;
    int rightheight = 0;

    if(!_isbalance(root->_left,leftheight) || !_isbalance(root->_right,rightheight))
    {
        return false;
    }

    if(abs(leftheight - rightheight) >= 2)
    {
        cout << root->_kv.first << "不平衡"  << endl;
        return false;
    }
    
    if(rightheight - leftheight != root->_bf)
    {
        cout << root->_kv.first << "平衡因子出现异常"  << endl;
        return false;
    }

    height = leftheight > rightheight ? leftheight + 1 : rightheight + 1;

    return true;
}

template <class K, class V>
int AVLTree<K,V>::_height(Node* root)
{
    if(root == nullptr)
    {
        return 0;
    }

    int leftheight = _height(root->_left);
    int rightheight = _height(root->_right);

    return leftheight > rightheight ? leftheight + 1 : rightheight + 1;
}

template <class K, class V>
void AVLTree<K,V>::_inorder(Node* root)
{
    if(root == nullptr)
    {
        return;
    }
    _inorder(root->_left);
    cout << root->_kv.first << "[" << root->_bf << "]" << endl;
    _inorder(root->_right);
}


template <class K, class V>
bool AVLTree<K,V>::insert(const pair<K,V>& kv)
{
    if(_root == nullptr)
    {
        _root = new Node(kv);
        return true;
    }
    else
    {
        Node* cur = _root;
        Node* parent = nullptr;
        while(cur)
        {
            if(cur->_kv.first < kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if(cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
                return false;
        }

        cur = new Node(kv);
        if(parent->_kv.first < kv.first)
        {
            parent->_right = cur;
        }
        else
        {
            parent->_left = cur;
        }
        cur->_parent = parent;

        while(parent)
        {
            if(cur == parent->_left)
            {
                parent->_bf--;
            }
            else
            {
                parent->_bf++;
            }

            if(parent->_bf == 0)
                break;
            else if(parent->_bf == 1 || parent->_bf == -1)
            {
                cur = cur->_parent;
                parent = parent->_parent;
            }
            else if(parent->_bf == 2 || parent->_bf == -2)
            {
                if(parent->_bf == 2 && cur->_bf == 1)
                {
                    RotateL(parent);
                }
                else if(parent->_bf == -2 && cur->_bf == -1)
                {
                    RotateR(parent);
                }
                else if(parent->_bf == -2 && cur->_bf == 1)
                {
                    RotateLR(parent);
                }
                else
                {
                    RotateRL(parent);
                }
                break;
            }
            else
                assert(false);//树有问题
        }

        return true;
    }
}

template <class K, class V>
void AVLTree<K,V>::RotateL(Node* parent)
{
    Node* subR = parent->_right;
    Node* subRL = subR->_left;
    Node* ppnode = parent->_parent;

    parent->_right = subRL;
    if(subRL)
        subRL->_parent = parent;

    subR->_left = parent;
    parent->_parent = subR;
    
    if(parent == _root)
    {
        _root = subR;
        subR->_parent = nullptr;
    }
    else
    {
        if(ppnode->_left == parent)
        {
            ppnode->_left = subR;
        }
        else
        {
            ppnode->_right = subR;
        }
        subR->_parent = ppnode;
    }

    parent->_bf = 0;
    subR->_bf = 0;
}

template <class K, class V>
void AVLTree<K,V>::RotateR(Node* parent)
{
    Node* subL = parent->_left;
    Node* subLR = subL->_right;
    Node* ppnode = parent->_parent;

    parent->_left = subLR;
    if(subLR)
        subLR->_parent = parent;
    
    subL->_right = parent;
    parent->_parent = subL;

    if(_root == parent)
    {
        _root = subL;
        subL->_parent = nullptr;
    }
    else
    {
        if(ppnode->_left == parent)
        {
            ppnode->_left = subL;
        }
        else
        {
            ppnode->_right = subL;
        }
        subL->_parent = ppnode;
    }

    parent->_bf = 0;
    subL->_bf = 0;
}

template <class K, class V>
void AVLTree<K,V>::RotateLR(Node* parent)
{
    Node* subL = parent->_left;
    Node* subLR = subL->_right;

    int bf = subLR->_bf;

    RotateL(parent->_left);
    RotateR(parent);

    if(bf == -1)
    {
        subLR->_bf = 0;
        subL->_bf = 0;
        parent->_bf = 1;
    }
    else if(bf == 1)
    {
        subLR->_bf = 0;
        subL->_bf = 1;
        parent->_bf = 0;
    }
    else if(bf == 0)
    {
        subLR->_bf = 0;
        subL->_bf = 0;
        parent->_bf = 0;
    }
    else
        assert(false);
}

template <class K, class V>
void AVLTree<K,V>::RotateRL(Node* parent)
{
    Node* subR = parent->_right;
    Node* subRL = subR->_left;

    int bf = subRL->_bf;
    RotateR(subR);
    RotateL(parent);

    if(bf == 1)
    {
        subR->_bf = 0;
        parent->_bf = -1;
        subRL->_bf = 0;
    }
    else if(bf == -1)
    {
        subR->_bf = 1;
        parent->_bf = 0;
        subRL->_bf = 0;
    }
    else if(bf == 0)
    {
        subR->_bf = 0;
        parent->_bf = 0;
        subRL->_bf = 0;
    }
    else
        assert(false);
}


